package cn.zwx.common.limiter.service;

import cn.zwx.base.enums.ResponseCodeEnum;
import cn.zwx.base.exception.BizWrapException;
import cn.zwx.common.limiter.annotation.LimitRule;
import cn.zwx.common.limiter.enums.LimiterTypeEnum;
import cn.zwx.common.limiter.enums.PeriodExecuteCodeEnum;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * @author zhangwenxue
 */
@Service
public class RedisLimiterServiceImpl implements LimiterService{

    public static final Logger logger = LoggerFactory.getLogger(RedisLimiterServiceImpl.class);
    /**
     * --ARGV[1]:qos,单位时间内最多请求次数
     * --ARGV[2]:单位限流窗口时间
     * 0：表示错误，比如可能是redis故障、过载
     * 1：允许
     * 2：允许但是当前窗口内已到达上限
     * 3：拒绝
     */
    static final String PERIOD_SCRIPT = "local limit = tonumber(ARGV[1])\n" +
            "--窗口即一个单位限流周期,这里用过期模拟窗口效果,等于p.permit\n" +
            "local window = tonumber(ARGV[2])\n" +
            "--请求次数+1,获取请求总数\n" +
            "local current = redis.call(\"INCRBY\", KEYS[1], 1)\n" +
            "--如果是第一次请求,则设置过期时间并返回 成功\n" +
            "if current == 1 then\n" +
            "    redis.call(\"expire\", KEYS[1], window)\n" +
            "end\n" +
            "--如果当前请求数量小于limit则返回 成功\n" +
            "if current < limit then\n" +
            "    return 1\n" +
            "--如果当前请求数量==limit则返回 最后一次请求\n" +
            "elseif current == limit then\n" +
            "    return 2\n" +
            "--请求数量>limit则返回 失败\n" +
            "else\n" +
            "    return 0\n" +
            "end";

    @Resource
    private RedisTemplate<String,Object> redisTemplate;



    @Override
    public String getType() {
        return LimiterTypeEnum.redis.getMsg();
    }


    @Override
    public boolean take(LimitRule limitRule) throws BizWrapException {
        // 1.执行 lua脚本
        String namespace = limitRule.namespace();
        String key = limitRule.key();
        String redisKey = namespace + ":" + key;
        List<String> keyList = Collections.singletonList(redisKey);
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptText(PERIOD_SCRIPT);
        redisScript.setResultType(Long.class);
        Long executeCode = redisTemplate.execute(redisScript, keyList,
                limitRule.quota(), limitRule.period());
        if (Objects.isNull(executeCode)){
            logger.error("redis限流错误 key:{}",redisKey);
            throw new BizWrapException(ResponseCodeEnum.PERIOD_ERROR);
        }
        // 2.解析执行结果
        List<Integer> allowedList = Arrays.asList(PeriodExecuteCodeEnum.INTERNAL_ALLOWED.getCode(),
                PeriodExecuteCodeEnum.INTERNAL_HIT_QUOTA.getCode());
        return allowedList.contains(executeCode.intValue());
    }
}
